import os
import jieba
import pickle
import numpy as np
import chardet
from collections import Counter
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.pipeline import make_pipeline
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.metrics import classification_report, accuracy_score
from sklearn.feature_selection import SelectKBest, chi2

class AdvancedChineseNovelClassifier:
    def __init__(self):
        self.model = None
        self.vectorizer = None
        self.label_encoder = {}
        self.reverse_label_encoder = {}
        self.feature_importance = {}
        self.class_distinctive_features = {}
        self.training_history = []
        
    def detect_encoding(self, file_path):
        """检测文件编码"""
        try:
            # 读取文件的一部分来检测编码
            with open(file_path, 'rb') as file:
                raw_data = file.read(10000)  # 读取前10000字节进行检测
                result = chardet.detect(raw_data)
                encoding = result['encoding']
                confidence = result['confidence']
                
                # 如果置信度低，默认使用常见的中文编码
                if confidence < 0.6:
                    # 尝试常见的中文编码
                    encodings_to_try = ['gbk', 'gb2312', 'gb18030', 'utf-8', 'big5']
                    for enc in encodings_to_try:
                        try:
                            with open(file_path, 'r', encoding=enc) as test_file:
                                test_file.read(1000)
                            return enc
                        except:
                            continue
                    return 'utf-8'  # 默认回退到utf-8
                
                return encoding if encoding else 'utf-8'
        except Exception as e:
            print(f"检测文件编码时出错: {e}")
            return 'utf-8'
    
    def read_file_with_encoding(self, file_path):
        """使用检测到的编码读取文件"""
        encoding = self.detect_encoding(file_path)
        max_retries = 3
        encodings_to_try = [encoding, 'gbk', 'gb2312', 'gb18030', 'utf-8', 'big5']
        
        for i, enc in enumerate(encodings_to_try):
            try:
                with open(file_path, 'r', encoding=enc) as file:
                    content = file.read()
                print(f"成功读取文件 {os.path.basename(file_path)}，使用编码: {enc}")
                return content
            except UnicodeDecodeError as e:
                if i == len(encodings_to_try) - 1:
                    print(f"所有编码尝试失败，文件: {file_path}，错误: {e}")
                    # 最后一次尝试，使用errors='ignore'忽略错误字符
                    try:
                        with open(file_path, 'r', encoding=enc, errors='ignore') as file:
                            content = file.read()
                        print(f"使用编码 {enc} 并忽略错误字符读取文件: {os.path.basename(file_path)}")
                        return content
                    except Exception as final_error:
                        print(f"最终读取失败，文件: {file_path}，错误: {final_error}")
                        return None
                else:
                    print(f"使用编码 {enc} 读取失败，尝试下一个编码...")
            except Exception as e:
                print(f"读取文件 {file_path} 时发生未知错误: {e}")
                if i == len(encodings_to_try) - 1:
                    return None
        
        return None
    
    def preprocess_text(self, text, use_stopwords=True):
        """中文文本预处理：分词和清洗"""
        if text is None:
            return ""
            
        # 使用jieba进行分词
        try:
            words = jieba.cut(text)
        except Exception as e:
            print(f"分词时出错: {e}")
            return ""
        
        # 基础停用词列表（可根据需要扩展）
        stopwords = set(['的', '了', '在', '是', '我', '有', '和', '就', '不', '人', '都', '一', '一个', '上', '也', '很', '到', '说', '要', '去', '你', '会', '着', '没有', '看', '自己', '这', '那', '他', '她', '它'])
        
        # 过滤处理
        if use_stopwords:
            filtered_words = [word.strip() for word in words 
                             if word.strip() and 
                             len(word) > 0 and
                             (len(word) > 1 or (len(word) == 1 and '\u4e00' <= word <= '\u9fff')) and
                             word not in stopwords]
        else:
            filtered_words = [word.strip() for word in words 
                             if word.strip() and 
                             len(word) > 0 and
                             (len(word) > 1 or (len(word) == 1 and '\u4e00' <= word <= '\u9fff'))]
        
        return ' '.join(filtered_words)
    
    def analyze_class_differences(self, texts, labels, feature_names, top_n=20):
        """分析不同类别之间的特征差异"""
        # 按类别分组文本
        class_texts = {}
        for text, label in zip(texts, labels):
            if label not in class_texts:
                class_texts[label] = []
            class_texts[label].append(text)
        
        # 计算每个类别的TF-IDF向量
        class_tfidf = {}
        for label, label_texts in class_texts.items():
            # 合并该类别的所有文本
            combined_text = ' '.join(label_texts)
            # 计算TF-IDF
            tfidf_vector = self.vectorizer.transform([combined_text])
            class_tfidf[label] = tfidf_vector.toarray()[0]
        
        # 找出每个类别最独特的特征
        distinctive_features = {}
        for label, tfidf_vector in class_tfidf.items():
            # 获取该类别中TF-IDF分数最高的特征
            top_indices = np.argsort(tfidf_vector)[-top_n:][::-1]
            top_features = [(feature_names[i], tfidf_vector[i]) for i in top_indices if tfidf_vector[i] > 0]
            distinctive_features[label] = top_features
        
        return distinctive_features
    
    def extract_style_features(self, text):
        """提取文本风格特征"""
        if not text:
            return {}
            
        words = text.split()
        if len(words) == 0:
            return {}
        
        # 计算各种文本特征
        total_chars = len(text.replace(' ', ''))
        total_words = len(words)
        avg_word_length = total_chars / total_words if total_words > 0 else 0
        
        # 词频统计
        word_freq = Counter(words)
        common_words = word_freq.most_common(10)
        
        # 段落特征（通过换行符识别）
        paragraphs = [p for p in text.split('\n') if p.strip()]
        avg_paragraph_length = sum(len(p.split()) for p in paragraphs) / len(paragraphs) if paragraphs else 0
        
        return {
            'avg_word_length': avg_word_length,
            'vocabulary_size': len(word_freq),
            'common_words': common_words,
            'paragraph_count': len(paragraphs),
            'avg_paragraph_length': avg_paragraph_length
        }
    
    def load_training_data(self, folder_path, novel_type):
        """加载训练数据并提取特征"""
        texts = []
        labels = []
        style_features = []
        
        if not os.path.exists(folder_path):
            print(f"错误：文件夹路径 '{folder_path}' 不存在")
            return texts, labels, style_features
            
        # 遍历文件夹中的所有txt文件
        txt_files = [f for f in os.listdir(folder_path) if f.endswith('.txt')]
        print(f"在文件夹 '{folder_path}' 中找到 {len(txt_files)} 个txt文件")
        
        for filename in txt_files:
            file_path = os.path.join(folder_path, filename)
            try:
                # 使用新的文件读取方法
                content = self.read_file_with_encoding(file_path)
                
                if content is None:
                    print(f"跳过无法读取的文件: {filename}")
                    continue
                
                # 预处理文本
                processed_text = self.preprocess_text(content)
                
                # 检查处理后的文本是否有效
                if not processed_text or len(processed_text.strip()) < 100:
                    print(f"警告: 文件 {filename} 内容过少或处理失败，跳过")
                    continue
                
                texts.append(processed_text)
                labels.append(novel_type)
                
                # 提取风格特征
                style_feat = self.extract_style_features(processed_text)
                style_features.append(style_feat)
                
            except Exception as e:
                print(f"处理文件 {filename} 时出错: {e}")
                continue
        
        print(f"从 '{folder_path}' 成功加载了 {len(texts)} 个 {novel_type} 类型的小说")
        return texts, labels, style_features
    
    def train_model(self, training_data_dict, model_type='naive_bayes', use_cross_validation=True):
        """
        训练分类模型
        training_data_dict: 字典，键为文件夹路径，值为小说类型
        model_type: 模型类型，可选 'naive_bayes', 'logistic_regression', 'random_forest'
        """
        all_texts = []
        all_labels = []
        all_style_features = []
        
        # 加载所有训练数据
        for folder_path, novel_type in training_data_dict.items():
            texts, labels, style_features = self.load_training_data(folder_path, novel_type)
            all_texts.extend(texts)
            all_labels.extend(labels)
            all_style_features.extend(style_features)
        
        if len(all_texts) == 0:
            print("错误：没有找到有效的训练数据")
            return False
        
        # 创建标签编码器
        unique_labels = list(set(all_labels))
        self.label_encoder = {label: idx for idx, label in enumerate(unique_labels)}
        self.reverse_label_encoder = {idx: label for label, idx in self.label_encoder.items()}
        
        # 将标签转换为数字
        y = np.array([self.label_encoder[label] for label in all_labels])
        
        # 创建TF-IDF向量器
        self.vectorizer = TfidfVectorizer(
            max_features=5000, 
            ngram_range=(1, 3),  # 扩展到3-gram
            min_df=2,            # 忽略出现次数太少的词
            max_df=0.8           # 忽略出现次数太多的词
        )
        
        # 转换文本为TF-IDF特征
        X_tfidf = self.vectorizer.fit_transform(all_texts)
        feature_names = self.vectorizer.get_feature_names_out()
        
        # 分析类别差异
        print("\n正在分析类别特征差异...")
        self.class_distinctive_features = self.analyze_class_differences(all_texts, all_labels, feature_names)
        
        # 打印每个类别的独特特征
        for label, features in self.class_distinctive_features.items():
            print(f"\n{label}类别的独特特征 (前10个):")
            for feature, score in features[:10]:
                print(f"  {feature}: {score:.4f}")
        
        # 特征选择
        print("\n进行特征选择...")
        k = min(2000, X_tfidf.shape[1])  # 选择最佳特征数量
        selector = SelectKBest(chi2, k=k)
        X_selected = selector.fit_transform(X_tfidf, y)
        
        # 获取选中的特征名称
        selected_indices = selector.get_support(indices=True)
        selected_features = [feature_names[i] for i in selected_indices]
        
        print(f"从 {X_tfidf.shape[1]} 个特征中选择了 {k} 个最相关的特征")
        
        # 选择模型
        if model_type == 'logistic_regression':
            model = LogisticRegression(
                C=1.0, 
                max_iter=1000,
                class_weight='balanced'
            )
        elif model_type == 'random_forest':
            model = RandomForestClassifier(
                n_estimators=100,
                max_depth=20,
                class_weight='balanced',
                random_state=42
            )
        else:  # 默认使用朴素贝叶斯
            model = MultinomialNB(alpha=0.1)
        
        # 交叉验证评估
        if use_cross_validation and len(set(all_labels)) > 1:
            print("\n进行交叉验证...")
            cv_scores = cross_val_score(model, X_selected, y, cv=min(5, len(set(y))))
            print(f"交叉验证准确率: {cv_scores.mean():.4f} (+/- {cv_scores.std() * 2:.4f})")
        
        # 划分训练集和测试集
        X_train, X_test, y_train, y_test = train_test_split(
            X_selected, y, test_size=0.2, random_state=42, stratify=y
        )
        
        # 训练模型
        print("\n训练最终模型...")
        model.fit(X_train, y_train)
        
        # 在测试集上评估
        y_pred = model.predict(X_test)
        accuracy = accuracy_score(y_test, y_pred)
        print(f"测试集准确率: {accuracy:.4f}")
        
        # 保存模型和相关信息
        self.model = model
        self.selector = selector
        self.feature_names = feature_names
        self.selected_features = selected_features
        
        # 记录训练历史
        training_record = {
            'model_type': model_type,
            'accuracy': accuracy,
            'classes': list(self.label_encoder.keys()),
            'sample_count': len(all_texts),
            'feature_count': k
        }
        self.training_history.append(training_record)
        
        print(f"\n模型训练完成！")
        print(f"支持的分类类型: {list(self.label_encoder.keys())}")
        print(f"总训练样本: {len(all_texts)}")
        print(f"使用特征数量: {k}")
        
        return True
    
    def predict_novel(self, novel_path, return_analysis=False):
        """预测小说的类型，可返回详细分析"""
        if self.model is None:
            print("错误：请先训练模型")
            return None
        
        if not os.path.exists(novel_path):
            print(f"错误：文件 '{novel_path}' 不存在")
            return None
        
        try:
            # 使用新的文件读取方法
            content = self.read_file_with_encoding(novel_path)
            
            if content is None:
                print(f"无法读取文件: {novel_path}")
                return None
            
            # 预处理文本
            processed_text = self.preprocess_text(content)
            
            # 检查处理后的文本是否有效
            if not processed_text or len(processed_text.strip()) < 50:
                print(f"警告: 文件内容过少或处理失败")
                return None
            
            # 提取风格特征
            style_features = self.extract_style_features(processed_text)
            
            # 转换为TF-IDF特征
            X_tfidf = self.vectorizer.transform([processed_text])
            
            # 特征选择
            X_selected = self.selector.transform(X_tfidf)
            
            # 预测
            prediction = self.model.predict(X_selected)[0]
            probability = self.model.predict_proba(X_selected)[0]
            
            # 获取预测结果
            predicted_label = self.reverse_label_encoder[prediction]
            confidence = probability[prediction]
            
            # 获取所有类别的概率
            class_probabilities = {}
            for idx, prob in enumerate(probability):
                class_name = self.reverse_label_encoder[idx]
                class_probabilities[class_name] = prob
            
            result = {
                'predicted_class': predicted_label,
                'confidence': confidence,
                'all_probabilities': class_probabilities,
                'style_features': style_features,
                'processed_text': processed_text
            }
            
            # 如果需要详细分析，计算特征贡献
            if return_analysis and hasattr(self.model, 'coef_'):
                feature_contributions = self.analyze_feature_contributions(
                    processed_text, predicted_label
                )
                result['feature_contributions'] = feature_contributions
            
            return result
            
        except Exception as e:
            print(f"预测时出错: {e}")
            return None
    
    def analyze_feature_contributions(self, processed_text, predicted_label):
        """分析特征对分类结果的贡献"""
        if not hasattr(self.model, 'coef_'):
            return {}
        
        # 获取特征权重
        if hasattr(self.model, 'coef_'):
            coef = self.model.coef_[self.label_encoder[predicted_label]]
        else:
            return {}
        
        # 获取文本中的特征
        words = processed_text.split()
        word_freq = Counter(words)
        
        # 计算每个特征的贡献
        contributions = {}
        for word in word_freq:
            if word in self.selected_features:
                idx = list(self.selected_features).index(word)
                contributions[word] = {
                    'frequency': word_freq[word],
                    'weight': coef[idx],
                    'contribution': word_freq[word] * coef[idx]
                }
        
        # 按贡献度排序
        sorted_contributions = sorted(
            contributions.items(), 
            key=lambda x: abs(x[1]['contribution']), 
            reverse=True
        )[:20]  # 取前20个最重要的特征
        
        return dict(sorted_contributions)
    
    def compare_with_training_data(self, novel_path):
        """将待分类小说与训练数据进行对比分析"""
        result = self.predict_novel(novel_path, return_analysis=True)
        if result is None:
            return
        
        print(f"\n=== 小说分类详细分析 ===")
        print(f"文件: {os.path.basename(novel_path)}")
        print(f"预测类型: {result['predicted_class']}")
        print(f"置信度: {result['confidence']:.4f}")
        
        print(f"\n所有类型概率:")
        for class_name, prob in sorted(result['all_probabilities'].items(), 
                                     key=lambda x: x[1], reverse=True):
            print(f"  {class_name}: {prob:.4f}")
        
        print(f"\n文本风格特征:")
        style = result['style_features']
        print(f"  平均词长: {style['avg_word_length']:.2f}")
        print(f"  词汇量: {style['vocabulary_size']}")
        print(f"  段落数: {style['paragraph_count']}")
        print(f"  平均段落长度: {style['avg_paragraph_length']:.2f}")
        
        if 'feature_contributions' in result:
            print(f"\n关键分类特征 (前10个):")
            contributions = result['feature_contributions']
            for i, (word, info) in enumerate(list(contributions.items())[:10]):
                print(f"  {i+1}. {word}: 频率={info['frequency']}, 权重={info['weight']:.4f}")
        
        # 与训练数据对比
        predicted_class = result['predicted_class']
        if predicted_class in self.class_distinctive_features:
            print(f"\n与{predicted_class}训练数据的特征对比:")
            training_features = [feat[0] for feat in self.class_distinctive_features[predicted_class][:10]]
            # 从 predict_novel 的返回结果中获取预处理文本
            processed_text = result.get('processed_text', '')
            novel_words = set(processed_text.split()) if processed_text else set()
            common_features = set(training_features) & novel_words
            
            print(f"  与训练数据共享的关键特征: {len(common_features)}个")
            if common_features:
                print(f"  共享特征示例: {list(common_features)[:5]}")
    
    def save_model(self, model_path):
        """保存训练好的模型"""
        if self.model is None:
            print("错误：没有训练好的模型可以保存")
            return False
        
        try:
            with open(model_path, 'wb') as f:
                pickle.dump({
                    'model': self.model,
                    'vectorizer': self.vectorizer,
                    'selector': self.selector,
                    'label_encoder': self.label_encoder,
                    'reverse_label_encoder': self.reverse_label_encoder,
                    'feature_names': self.feature_names,
                    'selected_features': self.selected_features,
                    'class_distinctive_features': self.class_distinctive_features,
                    'training_history': self.training_history
                }, f)
            print(f"模型已保存到: {model_path}")
            return True
        except Exception as e:
            print(f"保存模型时出错: {e}")
            return False
    
    def load_model(self, model_path):
        """加载已训练的模型"""
        try:
            with open(model_path, 'rb') as f:
                data = pickle.load(f)
            
            self.model = data['model']
            self.vectorizer = data['vectorizer']
            self.selector = data['selector']
            self.label_encoder = data['label_encoder']
            self.reverse_label_encoder = data['reverse_label_encoder']
            self.feature_names = data['feature_names']
            self.selected_features = data['selected_features']
            self.class_distinctive_features = data['class_distinctive_features']
            self.training_history = data['training_history']
            
            print(f"模型已从 {model_path} 加载")
            print(f"支持的分类类型: {list(self.label_encoder.keys())}")
            print(f"训练历史: {len(self.training_history)} 次训练记录")
            return True
        except Exception as e:
            print(f"加载模型时出错: {e}")
            return False

def main():
    classifier = AdvancedChineseNovelClassifier()
    
    while True:
        print("\n=== 高级中文小说分类系统 ===")
        print("1. 训练新模型")
        print("2. 加载现有模型")
        print("3. 对小说进行分类")
        print("4. 详细分析小说")
        print("5. 保存模型")
        print("6. 显示训练历史")
        print("7. 退出")
        
        choice = input("请选择操作 (1-7): ").strip()
        
        if choice == '1':
            print("\n--- 训练新模型 ---")
            training_data = {}
            
            while True:
                folder_path = input("请输入训练数据文件夹路径 (输入 'done' 结束添加): ").strip()
                if folder_path.lower() == 'done':
                    break
                
                novel_type = input("请输入该文件夹中小说的类型: ").strip()
                training_data[folder_path] = novel_type
            
            if training_data:
                print("\n选择模型类型:")
                print("1. 朴素贝叶斯 (默认)")
                print("2. 逻辑回归")
                print("3. 随机森林")
                model_choice = input("请选择 (1-3): ").strip()
                
                model_type = 'naive_bayes'
                if model_choice == '2':
                    model_type = 'logistic_regression'
                elif model_choice == '3':
                    model_type = 'random_forest'
                
                success = classifier.train_model(training_data, model_type=model_type)
                if success:
                    print("模型训练成功！")
            else:
                print("未提供训练数据")
        
        elif choice == '2':
            print("\n--- 加载现有模型 ---")
            model_path = input("请输入模型文件路径: ").strip()
            classifier.load_model(model_path)
        
        elif choice == '3':
            print("\n--- 小说分类 ---")
            if classifier.model is None:
                print("请先训练或加载模型")
                continue
            
            novel_path = input("请输入要分类的小说文件路径: ").strip()
            result = classifier.predict_novel(novel_path)
            
            if result:
                print(f"\n分类结果:")
                print(f"最可能的小说类型: {result['predicted_class']}")
                print(f"置信度: {result['confidence']:.4f}")
        
        elif choice == '4':
            print("\n--- 详细分析小说 ---")
            if classifier.model is None:
                print("请先训练或加载模型")
                continue
            
            novel_path = input("请输入要分析的小说文件路径: ").strip()
            classifier.compare_with_training_data(novel_path)
        
        elif choice == '5':
            print("\n--- 保存模型 ---")
            if classifier.model is None:
                print("没有训练好的模型可以保存")
                continue
            
            model_path = input("请输入模型保存路径: ").strip()
            classifier.save_model(model_path)
        
        elif choice == '6':
            print("\n--- 训练历史 ---")
            if not classifier.training_history:
                print("没有训练历史记录")
            else:
                for i, record in enumerate(classifier.training_history):
                    print(f"\n训练记录 {i+1}:")
                    print(f"  模型类型: {record['model_type']}")
                    print(f"  准确率: {record['accuracy']:.4f}")
                    print(f"  分类类型: {record['classes']}")
                    print(f"  样本数量: {record['sample_count']}")
                    print(f"  特征数量: {record['feature_count']}")
        
        elif choice == '7':
            print("再见！")
            break
        
        else:
            print("无效选择，请重新输入")

if __name__ == "__main__":
    # 安装依赖的提示
    print("请确保已安装以下依赖:")
    print("pip install jieba scikit-learn chardet")
    
    main()